package com.abc.ollama.controller;

import com.abc.ollama.entity.Embedding;
import com.abc.ollama.service.impl.OllamaClientService;
import com.abc.ollama.vo.OllamaChatRequest;
import com.abc.ollama.vo.OllamaEmbeddingRequest;
import com.abc.ollama.vo.OllamaEmbeddingResponse;
import com.abc.ollama.vo.OllamaGenerateRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.File;
import java.io.IOException;
import java.util.List;

@Slf4j
@RestController
@RequestMapping("ollama")
public class OllamaController {

    private final OllamaClientService ollamaClientService;

    public OllamaController(OllamaClientService ollamaClientService) {
        this.ollamaClientService = ollamaClientService;
    }


    @GetMapping("/generate/stream")
    public Flux<String> streamGenerate(@RequestParam String msg) {
        OllamaGenerateRequest ollamaGenerateRequest = new OllamaGenerateRequest();
        ollamaGenerateRequest.setPrompt(msg);
        return ollamaClientService.streamGenerate(ollamaGenerateRequest);
    }

    @GetMapping("/generate/stream/json")
    public Flux<String> streamGenerateJson(@RequestParam String msg) {
        OllamaGenerateRequest ollamaGenerateRequest = new OllamaGenerateRequest();
        ollamaGenerateRequest.setPrompt(msg + ". Response using json");
        ollamaGenerateRequest.setFormat("json");
        return ollamaClientService.streamGenerate(ollamaGenerateRequest);
    }

    @PostMapping("/chat/stream")
    public Flux<String> chatStream(@RequestBody OllamaChatRequest chatRequest) {
        return ollamaClientService.streamChatResponse(chatRequest);
    }

    @PostMapping("/chat/string")
    public String chatString(@RequestBody OllamaChatRequest chatRequest) {
        return ollamaClientService.stringChatResponse(chatRequest);
    }

    @PostMapping("/embeddings")
    public OllamaEmbeddingResponse embeddings(@RequestBody OllamaEmbeddingRequest ollamaEmbeddingRequest) {
        return ollamaClientService.embeddings(ollamaEmbeddingRequest);
    }

    @PostMapping("/embeddings/save")
    public void saveEmbeddings(@RequestBody OllamaEmbeddingRequest ollamaEmbeddingRequest) {
        // text-embed + 时间戳
        String fileName = "text-embed-" + System.currentTimeMillis();
        ollamaClientService.saveEmbeddings(ollamaEmbeddingRequest, fileName);
    }

    // localKnowledgeChat
    @PostMapping("/chat/localKnowledge")
    public Flux<String> localKnowledgeChat(@RequestBody OllamaChatRequest chatRequest) {
        return ollamaClientService.localKnowledgeChat(chatRequest);
    }

    // embeddingsFile
    @GetMapping("/embeddings/file")
    public void embeddingsFile(@RequestParam String filePath) {
        ollamaClientService.embeddingsFile(filePath, "test");
    }

    // 将上面那个改成直接上传文件
    @PostMapping("/embeddings/file/upload")
    public void embeddingsFileUpload(@RequestParam("file") MultipartFile file) {
        try {
            String filePath = "D:\\embeddings\\" + file.getOriginalFilename();
            file.transferTo(new File(filePath));
            log.info("File uploaded successfully.");
            ollamaClientService.embeddingsFile(filePath, file.getOriginalFilename());
        } catch (IOException e) {
            log.error("Error uploading file: {}", e.getMessage());
        }
    }

    // 查询所有该目录下的文件
    @GetMapping("/embeddings/file/list")
    public List<Embedding> getAllFiles() {
        return ollamaClientService.getAllFiles();
    }

    // deleteEmbeddingsByFileId
    @GetMapping("/embeddings/file/delete")
    public void deleteEmbeddingsByFileId(@RequestParam String fileId) {
        ollamaClientService.deleteEmbeddingsByFileId(fileId);
    }

    // deepseek
    @PostMapping("/chat/deepseek")
    public Flux<String> deepseekChat(@RequestBody OllamaChatRequest chatRequest) {
        return ollamaClientService.deepseekChat(chatRequest);
    }

    // deepseekChatNotStream
    @PostMapping("/chat/deepseek/notStream")
    public String deepseekChatNotStream(@RequestBody OllamaChatRequest chatRequest) {
        return ollamaClientService.deepseekChatNotStream(chatRequest);
    }

    // get related doc
    @PostMapping("/chat/deepseek/getRelatedDoc")
    public List<Embedding> getRelatedDoc(@RequestBody OllamaChatRequest chatRequest) {
        String msg = chatRequest.getMessages().get(chatRequest.getMessages().size() - 1).getContent();
        return ollamaClientService.getRelatedDoc(msg);
    }
}
